(*
  Copyright (c) 2009 Barry Schwartz

  Permission is hereby granted, free of charge, to any person
  obtaining a copy of this software and associated documentation
  files (the "Software"), to deal in the Software without
  restriction, including without limitation the rights to use,
  copy, modify, merge, publish, distribute, sublicense, and/or sell
  copies of the Software, and to permit persons to whom the
  Software is furnished to do so, subject to the following
  conditions:

  The above copyright notice and this permission notice shall be
  included in all copies or substantial portions of the Software.

  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
  EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
  OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
  NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
  HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
  WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
  FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
  OTHER DEALINGS IN THE SOFTWARE.
*)

(*-----------------------------------------------------------------------*)
open UTypes
(*-----------------------------------------------------------------------*)

type 'a trie = {
  t_tree : int array; (* array of 3n entries: offset, char, data, offset, char, data,... *)
  t_data : 'a array;  (* array of the actual data *)
  mutable t_data_len : int; (* number of used enties in the |t_data| array *)
}

let make len = {
  t_tree = Array.make (3 * len) (-1);
  t_data = Array.make len (Obj.magic 0);
  t_data_len = 0;
}

let shrink trie len = {
  t_tree = Array.sub trie.t_tree 0 (3 * len);
  t_data = Array.sub trie.t_data 0 trie.t_data_len;
  t_data_len = trie.t_data_len;
}

let get_offset trie pos = trie.t_tree.(3 * pos)
let get_char trie pos = trie.t_tree.((3 * pos) + 1)
let get_data_index trie pos = trie.t_tree.((3 * pos) + 2)
let set_offset trie pos x = trie.t_tree.(3 * pos) <- x
let set_char trie pos x = trie.t_tree.((3 * pos) + 1) <- x
let set_data_index trie pos x = trie.t_tree.((3 * pos) + 2) <- x

let get_data trie pos =
  let index = get_data_index trie pos in
    if index < 0 then
      None
    else
      Some trie.t_data.(index)

let set_data trie pos x =
  let index = get_data_index trie pos in
    if index < 0 then
      begin
        trie.t_data.(trie.t_data_len) <- x;
        trie.t_tree.((3 * pos) + 2) <- trie.t_data_len;
        trie.t_data_len <- trie.t_data_len + 1
      end
    else
      trie.t_data.(index) <- x
  
(* |lookup <trie> <pos> <str>| checks whether <str> occures in <trie> starting at <pos>. *)
let lookup trie pos str =
  if (Array.length str) = 0 then
    None
  else
    let rec iter pos i =
      let chr = str.(i) in
        if (get_char trie (pos + chr)) <> chr then
          None
        else if (i + 1) >= (Array.length str) then
          get_data trie (pos + chr)
        else
          iter (get_offset trie (pos + chr)) (i + 1)
    in
      iter pos 0
  
let lookup_list trie pos str =
  let rec iter pos chr str =
    if (get_char trie (pos + chr)) <> chr then
      None
    else
      begin
        match str with
          | [] -> get_data trie (pos + chr)
          | c :: cs -> iter (get_offset trie (pos + chr)) c cs
      end
  in
    match str with
      | [] -> None
      | c :: cs -> iter pos c cs
  
(*
  |lookup_prefix <trie> <pos> <str> <start>| checks whether a prefix of the substring of <str> that
  starts at <start> occures in <trie> starting at <pos>.
*)
let lookup_prefix trie pos str start =
  if start >= Array.length str then
    None
  else
    let rec iter pos result i =
      let chr = str.(i) in
        if (get_char trie (pos + chr)) <> chr then
          result
        else
          let new_result =
            match get_data trie (pos + chr) with
              | None -> result
              | x -> x
          in
            if (i + 1) >= (Array.length str) then
              new_result
            else
              iter (get_offset trie (pos + chr)) new_result (i + 1)
    in
      iter pos None start
  
let lookup_prefix_list trie pos str =
  let rec iter pos result chr str =
    if (get_char trie (pos + chr)) <> chr then
      result
    else
      let new_result =
        match get_data trie (pos + chr) with | None -> result | x -> x
      in
        match str with
          | [] -> new_result
          | c :: cs -> iter (get_offset trie (pos + chr)) new_result c cs
  in
    match str with
      | [] -> None
      | c :: cs -> iter pos None c cs

(*-----------------------------------------------------------------------*)
  
module SimpleTrie =

struct
  type 'a node = {
    tn_offset : int;
    tn_char   : uc_char;
    tn_data   : 'a option
  }
      
  type 'a simple_trie = (('a node) list) array
      
  let make size = Array.make size []
    
  let next trie pos char =
    let rec iter children =
      match children with
        | [] -> raise Not_found
        | c :: cs -> if c.tn_char = char then c.tn_offset else iter cs
    in
      iter trie.(pos)
        
  let add_child trie pos char index =
    let child =
      { tn_char = char;
        tn_offset = index;
        tn_data = None }
    in
      trie.(pos) <-
        let rec iter children =
          match children with
            | [] -> [ child ]
            | c :: cs ->
                if c.tn_char < char then
                  c :: (iter cs)
                else if c.tn_char = char then
                  child :: cs
                else
                  child :: children
        in
          iter trie.(pos)
            
  let set trie pos char data =
    trie.(pos) <-
      let rec iter children =
        match children with
          | [] -> raise Not_found
          | c :: cs ->
              if c.tn_char = char then
                { c with tn_data = Some data } :: cs
              else
                c :: (iter cs)
      in
        iter trie.(pos)

  (* |build <data>| creates a simple_trie from a list of (key, value) pairs. *)
  let build data =
    let rec calc_size size list =
      match list with
        | [] -> size
        | (str, _) :: tail -> calc_size (size + (Array.length str)) tail
    in
    let size = calc_size 1 data in
    let trie = make size in
    let rec fill_trie free data =
      let rec insert pos free i str data =
        if i >= (Array.length str) then
          free
        else if i + 1 = Array.length str then
          begin
            try
              (set trie pos str.(i) data; free)
            with
              | Not_found ->
                  add_child trie pos str.(i) free;
                  set trie pos str.(i) data;
                  free + 1
          end
        else
          begin
            try
              insert (next trie pos str.(i)) free (i + 1) str data
            with
              | Not_found ->
                  add_child trie pos str.(i) free;
                  insert free (free + 1) (i + 1) str data
          end
      in
        match data with
          | [] -> Array.sub trie 0 free
          | (s, d) :: tail -> fill_trie (insert 0 free 0 s d) tail
    in
      fill_trie 1 data
        
  (* |compress <trie>| translates a simple_trie into a trie. *)
  let compress trie =
    let len = Array.length trie in
    let ht = Hashtbl.create len in
      (* merge duplicates *)
      for i = len - 1 downto 0 do
        trie.(i) <-
          let rec iter children =
            match children with
              | [] -> []
              | c :: cs ->
                  try
                    let x = Hashtbl.find ht trie.(c.tn_offset) in
                      { (c) with tn_offset = x; } :: (iter cs)
                  with
                    | Not_found ->
                        Hashtbl.add ht trie.(c.tn_offset) c.tn_offset;
                        c :: (iter cs)
          in
            iter trie.(i)
      done
end                                     (* module SimpleTrie *)

(*-----------------------------------------------------------------------*)

let build max_code data =
  let simple_trie = SimpleTrie.build data in
  let len = (max_code + 1) * (Array.length simple_trie) in
    SimpleTrie.compress simple_trie;
    let trie = make len in
    let used = Array.make len false in
    let alloc_node children =
      let rec find pos children =
        if used.(pos) then
          find (pos + 1) children
        else
          let rec iter ch =
            match ch with
              | [] -> pos
              | c :: cs ->
                  if (get_char trie (pos + c.SimpleTrie.tn_char)) < 0 then
                    iter cs
                  else
                    find (pos + 1) children
          in
            iter children
      in
      let pos = find 0 children in
        used.(pos) <- true;
        List.iter
          (fun c -> set_char trie (pos + c.SimpleTrie.tn_char) c.SimpleTrie.tn_char)
          children;
        pos
    in
      
    (* mapping of indices: simple_trie to trie *)
    let alloc_table = Array.make (Array.length simple_trie) (-1) in
    let rec alloc pos =
      if alloc_table.(pos) < 0 then
        begin
          alloc_table.(pos) <- alloc_node simple_trie.(pos);
          List.iter (fun c -> alloc c.SimpleTrie.tn_offset) simple_trie.(pos)
        end
      else
        ()
    in
      alloc 0;
      (* mapping: data to index *)
      let ht = Hashtbl.create 2000 in
        begin
          let rec fill_trie pos =
            let trie_pos = alloc_table.(pos) in
            let children = simple_trie.(pos) in
              begin

                simple_trie.(pos) <- [];

                (* prevent writing this entry multiple times *)
                let rec iter children =
                  match children with
                    | [] -> ()
                    | c :: cs ->
                        set_offset trie (trie_pos + c.SimpleTrie.tn_char)
                          alloc_table.(c.SimpleTrie.tn_offset);
                        set_char trie (trie_pos + c.SimpleTrie.tn_char) c.SimpleTrie.tn_char;
                        begin
                          match c.SimpleTrie.tn_data with
                            | None -> ()
                            | Some x ->
                                try
                                  (* check whether |x| is already somewhere in the data array *)
                                  let i = Hashtbl.find ht x in
                                    set_data_index trie (trie_pos + c.SimpleTrie.tn_char) i
                                with
                                  | Not_found ->
                                      set_data trie (trie_pos + c.SimpleTrie.tn_char) x;
                                      Hashtbl.add ht x
                                        (get_data_index trie (trie_pos + c.SimpleTrie.tn_char))
                        end;
                        fill_trie c.SimpleTrie.tn_offset;
                        iter cs
                in
                  iter children
              end
          in
            fill_trie 0
        end;

        (* remove empty space at the end of the array *)
        let rec strip i =
          if (not used.(i)) && i > 0 then
            strip (i - 1)
          else
            shrink trie (i + max_code)
        in
          strip (len - 1)
        
(*-----------------------------------------------------------------------*)
